import librosa
import cv2
from scipy.io import wavfile
import os
import numpy as np
import soundfile as sf
import shutil


def data_augmentation(input_audio, output_audio):            # 时间拉伸
    in_path = os.path.basename(input_audio)
    out_path = os.path.dirname(output_audio)
    if not os.path.isdir(out_path):  # 判断目标路径是否存在，不存在则创建
        os.mkdir(out_path)

    shutil.copy(input_audio, out_path)

    y, sr = librosa.load(input_audio, sr=None)
    for rate in [0.93, 1.07]:  # 拉伸速率
        y_shift = librosa.effects.time_stretch(y, rate=rate)        # 使用TS生成新数据
        sf.write(os.path.join(out_path,
            in_path.replace('.wav', '_ts-{}.wav'.format(rate))), y_shift, sr)
    for step1 in [-1, 1]:
        y_shift = librosa.effects.pitch_shift(y, sr, n_steps=step1)          # 使用PS1生成新数据
        sf.write(os.path.join(out_path,
            in_path.replace('.wav', '_ps1-{}.wav'.format(step1))), y_shift, sr)
    for step2 in [-2.5, 2.5]:
        y_shift = librosa.effects.pitch_shift(y, sr, n_steps=step2)              # 使用PS2生成新数据
        sf.write(os.path.join(out_path,
            in_path.replace('.wav', '_ps2-{}.wav'.format(step2))), y_shift, sr)
    for step3 in [2, 3]:
        wn = np.random.normal(0, 1, len(y))                         # 从高斯分布（正态分布）提取样本
        y_noise = np.where(y != 0.0, y.astype('float64') + 0.02 * wn, 0.0).astype(np.float32)
        sf.write(os.path.join(out_path,
            in_path.replace('.wav', "_no-{}.wav".format(step3))),  y_noise, sr)
    


index = 0

save_path = "enhanced/"               # 增强后数据(输出)文件夹路径[无文件夹会自动创建]

g = os.walk('clustered')

for path, dir_list, file_list in g:
    for filename in file_list:
        print(filename)
        in_path = path+'/' + filename
        out_path = save_path+path[10:]+'/'
        print(out_path)
        if (len([lists for lists in os.listdir(path) if os.path.isfile(os.path.join(path, lists))])<5):
            data_augmentation(in_path, out_path)         # 数据增强 4*2
        else:
            if not os.path.isdir(out_path):  # 判断目标路径是否存在，不存在则创建
                os.mkdir(out_path)
            shutil.copy(in_path, out_path)

        index = index + 1
        print(index)